import os

from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough


def ZPAi_LLM():
    return ChatZhipuAI(
        temperature=0.5,
        model="glm-4",
        api_key=os.getenv('zp_ai'),
        base_url="https://open.bigmodel.cn/api/paas/v4/chat/completions"
    )


llm = ZPAi_LLM()


def GetNotTypePromptByAI(title, content):
    prompt = PromptTemplate.from_template("""
        你是一个有丰富笔记书写经验的人，请根据下面的文章标题和内容对文章进行分类，输出一个最多6个汉字的分类
        标题：{question}
        内容: {content}
        """)
    chain = {"question": RunnablePassthrough(), "content": RunnablePassthrough()} | prompt | llm | StrOutputParser()
    return chain.invoke({"question": title, "content": content})
